import torch
from torch import nn, optim
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

# 设定一些超参数
num_clients = 100
num_epochs = 5
batch_size = 32
num_rounds = 10
learning_rate = 0.01

# 检查GPU是否可用
device = "cuda" if torch.cuda.is_available() else "cpu"


# 定义模型架构
class NeuralNet(nn.Module):
    def __init__(self):
        super(NeuralNet, self).__init__()
        self.flatten = nn.Flatten()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.ReLU(),
            nn.Linear(128, 10),
        )

    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits


# 加载MNIST数据集
def load_data():
    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
    return train_dataset, test_dataset


# 训练客户端模型
def train_client(model, train_loader, epochs):
    model.train()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=learning_rate)
    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()


# 聚合模型权重
def aggregate_weights(global_model, client_models):
    global_dict = global_model.state_dict()
    for k in global_dict.keys():
        global_dict[k] = torch.stack([client_models[i].state_dict()[k].float() for i in range(len(client_models))],
                                     0).mean(0)
    global_model.load_state_dict(global_dict)
    for model in client_models:
        model.load_state_dict(global_model.state_dict())


# 联邦学习过程
def federated_learning(train_dataset, num_clients, num_epochs, batch_size, num_rounds):
    # 初始全局模型
    global_model = NeuralNet().to(device)

    # 划分数据集
    client_datasets = torch.utils.data.random_split(train_dataset,
                                                    [int(train_dataset.data.shape[0] / num_clients) for _ in
                                                     range(num_clients)])

    # 训练和聚合
    for round in range(num_rounds):
        client_models = [NeuralNet().to(device) for _ in range(num_clients)]
        for client_model in client_models:
            client_model.load_state_dict(global_model.state_dict())

        for i, client_model in enumerate(client_models):
            train_loader = DataLoader(client_datasets[i], batch_size=batch_size, shuffle=True)
            train_client(client_model, train_loader, num_epochs)

        aggregate_weights(global_model, client_models)

    return global_model


# 运行联邦学习
train_dataset, test_dataset = load_data()
global_model = federated_learning(train_dataset, num_clients, num_epochs, batch_size, num_rounds)


# 测试全局模型性能
def test_model(model, test_dataset):
    test_loader = DataLoader(test_dataset, batch_size=1000, shuffle=False)
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += nn.functional.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    print(
        f"Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.2f}%)")


test_model(global_model, test_dataset)
